# -*- coding: utf-8 -*-
import uuid
from typing import Container, Optional, Type, Dict
from pydantic import BaseModel, create_model, ConfigDict
from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy.orm import Session as BaseSession, RelationshipProperty
from sqlalchemy.inspection import inspect as sa_inspect


TYPE_MAPPING = {
    "UUID": uuid.UUID,
    "JSON": list,
    "GUID": uuid.UUID,
    "BINARY": uuid.UUID,
    # ...
}


class OrmConfig:
    model_config = ConfigDict(from_attributes=True)


def sqlalchemy_to_pydantic(
        db_model: Type,
        *,
        exclude: Container[str] = (),
        include: Container[str] = None
) -> Type[BaseModel]:
    mapper = sa_inspect(db_model)
    fields = {}

    # 包含所有属性，不仅是列
    for name, attr in mapper.all_orm_descriptors.items():
        if name.startswith('_') or name in exclude or (include and name not in include):
            continue

        # 处理列属性
        if hasattr(attr, 'property') and hasattr(attr.property, 'columns'):
            column = attr.property.columns[0]

            # 增强类型推断
            python_type = None
            try:
                if hasattr(column.type, "python_type"):
                    python_type = column.type.python_type
                elif hasattr(column.type, "impl") and hasattr(column.type.impl, "python_type"):
                    python_type = column.type.impl.python_type

                if str(column.type).split('(')[0] in TYPE_MAPPING.keys():
                    python_type = TYPE_MAPPING.get(str(column.type).split('(')[0], str)
            except (NotImplementedError, AttributeError):
                try:
                    python_type = TYPE_MAPPING.get(str(column.type).split('(')[0], str)
                except:
                    python_type = str

            # 处理Optional
            python_type = Optional[python_type] if column.nullable else python_type
            default = None if column.nullable else ...

            fields[name] = (python_type, default)

        # 处理混合属性
        elif isinstance(attr, hybrid_property):
            fields[name] = (attr.fget.__annotations__.get('return', str), None)

        # 处理关系字段
        elif hasattr(attr, 'property') and isinstance(attr.property, RelationshipProperty):
            continue  # 或者处理关系

    # 确保包含token字段（如果被意外排除）
    if 'token' not in fields and hasattr(db_model, 'token'):
        fields['token'] = (str, None)

    return create_model(
        f"{db_model.__name__}Model",
        __config__=ConfigDict(from_attributes=True),
        **fields
    )


class PydanticSession(BaseSession):
    """
    自动将查询结果转换为Pydantic模型的Session
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._pydantic_cache: Dict[Type, Type[BaseModel]] = {}

    def _get_pydantic_model(self, db_model: Type) -> Type[BaseModel]:
        """获取或创建Pydantic模型"""
        if db_model not in self._pydantic_cache:
            self._pydantic_cache[db_model] = sqlalchemy_to_pydantic(db_model)
        return self._pydantic_cache[db_model]

    def query(self, *entities, **kwargs):
        query = super().query(*entities, **kwargs)

        # 只包装返回单个实体的查询
        if len(entities) == 1:
            db_model = entities[0]
            pydantic_model = self._get_pydantic_model(db_model)

            def _convert_result(result):
                if result is None:
                    return None
                # 先将SQLAlchemy对象转换为字典
                # data = {c.key: getattr(result, c.key) for c in sa_inspect(result).mapper.columns}
                return pydantic_model.from_orm(result)

            # 包装常用查询方法
            original_first = query.first
            original_all = query.all
            original_one = query.one
            original_one_or_none = query.one_or_none

            def first_wrapper(*args, **kwargs):
                return _convert_result(original_first(*args, **kwargs))

            def all_wrapper(*args, **kwargs):
                return [_convert_result(r) for r in original_all(*args, **kwargs)]

            def one_wrapper(*args, **kwargs):
                return _convert_result(original_one(*args, **kwargs))

            def one_or_none_wrapper(*args, **kwargs):
                return _convert_result(original_one_or_none(*args, **kwargs))

            query.first = first_wrapper
            query.all = all_wrapper
            query.one = one_wrapper
            query.one_or_none = one_or_none_wrapper

        return query

    def add(self, instance, _warn=True):
        """重写add方法"""
        super().add(instance, _warn=_warn)
        return instance  # 保持链式调用

    def commit(self):
        """重写commit方法"""
        super().commit()
        # 提交后自动刷新所有待刷新对象
        for instance in self.new:
            self.refresh(instance)
        return self

    def flush(self, objects=None):
        """重写flush方法"""
        super().flush(objects)
        if objects:
            for instance in objects:
                self.refresh(instance)
        return self

    def _auto_convert(self, instance):
        """自动转换实例为Pydantic模型"""
        if instance is None:
            return None
        model = self._get_pydantic_model(instance.__class__)
        return model.model_validate(instance)

    def save(self, instance, commit=True):
        """新增的save便捷方法"""
        self.add(instance)
        if commit:
            self.commit()
        return self._auto_convert(instance)

# 使用示例
# if __name__ == "__main__":
#     from sqlalchemy import create_engine, Column, Integer, String
#     from sqlalchemy.ext.declarative import declarative_base
#     from sqlalchemy.orm import sessionmaker
#
#     Base = declarative_base()
#
#
#     class User(Base):
#         __tablename__ = 'users'
#         id = Column(Integer, primary_key=True)
#         name = Column(String)
#         age = Column(Integer)
#
#
#     engine = create_engine('sqlite:///:memory:')
#     Base.metadata.create_all(engine)
#
#     Session = sessionmaker(bind=engine, class_=PydanticSession)
#     session = Session()
#
#     # 添加测试数据
#     session.add(User(name="Alice", age=25))
#     session.add(User(name="Bob", age=30))
#     session.commit()
#
#     # 查询会自动返回Pydantic模型
#     user = session.query(User).first()
#     print(type(user))  # <class '__main__.UserModel'>
#     print(user.dict())  # {'id': 1, 'name': 'Alice', 'age': 25}
#
#     users = session.query(User).all()
#     print(type(users[0]))  # <class '__main__.UserModel'>
